#!/usr/bin/env python

import argparse
import psycopg2
import sys
import sblab
import pysam
import md5

parser = argparse.ArgumentParser(description= """

DESCRIPTION:
    Execute samtools idxstats on the bamfiles passed to --bams and import output
    to sblab.main.samtools_idxstats (output also printed to stdout).
    
    Bam files should be already index.
    
    Identical rows are imported only once or not imported if
    already found in samtools_idxstats. Distinct rows sharing the same
    primary key will cause exception.
    
EXAMPLES:
    ## Import all bam files in current dir:
    ls *.bam | idxstats_to_sblab.py -b -

REQUIRES:
    import argparse
    import psycopg2
    import sblab
    import pysam
    
TODO:
      
""", formatter_class= argparse.RawTextHelpFormatter)

parser.add_argument('-b', '--bams',
                    nargs='+',
                    type= str,
                    required= True,
                    help="""Input bam files, already indexed, for which samtools
idxstats should be exported. Use - to read the list of files from stdin.
                    """)

parser.add_argument('--no_md5',
                    action= 'store_true',
                    help="""Do not compute md5sum for bam files.
                    """)

parser.add_argument('--nocommit',
                    action= 'store_true',
                    help="""Do not commit changes to database.
                    """)

args = parser.parse_args()

# -----------------------------------------------------------------------------

def sumfile(fobj):
    '''Returns an md5 hash for an object with read() method.'''
    m = md5.new()
    while True:
        d = fobj.read(8096)
        if not d:
            break
        m.update(d)
    return(m.hexdigest())


def md5sum(fname):
    '''Returns an md5 hash for file fname, or stdin if fname is "-".'''
    if fname == '-':
        ret = sumfile(sys.stdin)
    else:
        try:
            f = file(fname, 'rb')
        except:
            return 'Failed to open file'
        ret = sumfile(f)
        f.close()
    return(ret)
# -----------------------------------------------------------------------------

if args.bams == ['-']:
    args.bams= sys.stdin.readlines()
    args.bams= [x.strip() for x in args.bams]

conn= psycopg2.connect(sblab.get_psycopgpass())
conn.set_isolation_level(0)
cur= conn.cursor()

## Put here parsed indexes
sql= 'CREATE TEMP TABLE idxstats_tmp (LIKE samtools_idxstats);'
cur.execute(sql)

for bam in args.bams:
    """
    MEMO: The output of pysam.idxstats(bam) look like:
    ['chr1\t140\t52\t0\n', 'chr2\t122\t5927\t0\n', ..., '*\t0\t0\t0\n']
    """
    if not args.no_md5:
        checksum= md5sum(bam)
    else:
        checksum= None
    try:
        idxstats= pysam.idxstats(bam)
    except:
        print('Skipping %s: Cannot execute idxstats' %(bam))
        continue
    idxstats= [x.strip().split('\t') for x in idxstats]
    for i in range(0, len(idxstats)):
        seq_line=  idxstats[i]
        seq_line.append(bam)
        seq_line.append(checksum)
        sql= 'INSERT INTO idxstats_tmp VALUES (' + ', '.join(['%s']*len(seq_line)) + ');'
        cur.execute(sql, seq_line)
        print('\t'.join(seq_line))
## Insert rows not already present.
sql= 'INSERT INTO samtools_idxstats SELECT * FROM idxstats_tmp EXCEPT SELECT * FROM samtools_idxstats; DROP TABLE idxstats_tmp;'
cur.execute(sql)

cur.close()
if not args.nocommit:
    conn.commit()
conn.close()
sys.exit()
    